from src.model import *
from src.utils import *
from keras.optimizers import SGD
import numpy as np

# 尝试跑自己的小网络
if __name__ == '__main__':

    epochs = 200
    train_filename = r'H:\wangjianlian\data\formal_data\HER2\thrid_generation\train'
    batch_size = 50
    if_save_model = True

    selfModel = SelfModel()
    model = selfModel.build(input_shape = (256, 256, 3))
    # model = selfModel.build(input_shape = (256, 256, 1)) # 边缘图
    model.compile(optimizer=SGD(lr=0.000005, momentum=0.9, nesterov=True),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    model.summary()

    data = Data()
    max_metrics = 0
    for epoch in range(epochs):
        print("************ 第%d代 *************"%(epoch + 1))
        all_loss = 0
        all_metrics = 0
        for index, trained_imgs, trained_label in data.read_image(train_filename, batch_size, shuffle = True):
            x_train = trained_imgs
            # x_train = np.expand_dims(x_train, axis=3)  # 增维，仅对边缘图
            y_train = trained_label[0]
            loss_and_metrics  = model.train_on_batch(x_train, y_train) # 应该是返回损失值和metrics
            print("训练第%d批，loss值:%f" % (index, loss_and_metrics[0]))
            print("训练第%d批，metrics值:%f" % (index, loss_and_metrics[1]))
            all_loss += loss_and_metrics[0]
            all_metrics += loss_and_metrics[1]

        print("************ 训练：平均值 *************")
        print("训练平均loss值:%f" % (all_loss / index))
        print("训练平均metrics值:%f" % (all_metrics / index))

        if if_save_model:
            if all_metrics/index > max_metrics:
                max_metrics = all_metrics / index
                model.save(r"H:\wangjianlian\project\Python\networkTest\resources\weight\temp\self_model" + "//" + str(epoch + 1) + 'model_1.h5')
